# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import math
import random

from coin.base.timestamp import get_timestamp
from coin.experimental.xguo.fruit.util.random_process import GBMotion
from coin.experimental.xguo.fruit.util.book_builder import BookBuilder
from coin.experimental.xguo.fruit.base.single_product import SingleProductPriceCalculator
from coin.experimental.xguo.fruit.util.util import VolatilityCalculator


class GbmBookGenerator(object):
  def __init__(self):
    current_t = get_timestamp() / 1e9
    gbm = GBMotion.create(mu=0.000002, sigma=0.0002, dt=0.5, s=8000, t=current_t)
    self._gbm = gbm
    self._book = BookBuilder('BTC-USDT')
    self._qty = 99999999
    self._tick = 1

  def __iter__(self):
    return self

  def __next__(self):
    value = next(self._gbm)
    v = value.s
    t = int(value.t * 1e9)
    bids = [[v - 0.5 * self._tick, self._qty]]
    asks = [[v + 0.5 * self._tick, self._qty]]
    self._book.snapshot(bids, asks, t)
    return self._book


class SineBookGenerator(object):
  def __init__(self, fsamp):
    self._current_t = get_timestamp()
    self._delta_t = int(5e8)
    self._book = BookBuilder('BTC-USDT')
    self._qty = 1e9
    self._tick = 0.5
    self._offset = 1000
    self._fsamp = fsamp

  def __iter__(self):
    return self

  def __next__(self):
    self._current_t += self._delta_t

    v = self._offset
    for fs, amp in self._fsamp:
      x = 2 * math.pi * fs * self._current_t
      v += self._offset + amp * math.sin(x)

    asks = []
    bids = []
    for n in range(1, 10, 2):
      nn = n / 2.
      asks += [v + nn * self._tick, self._qty]
      bids += [v - nn * self._tick, self._qty]
    self._book.snapshot(bids, asks, self._current_t)
    return self._book


class RandomWalkBookGenerator(object):
  def __init__(self):
    self._current_t = get_timestamp()
    self._current_price = 10000
    self._delta_t = int(5e8)
    self._book = BookBuilder('BTC-USDT')
    self._qty = 1e9
    self._tick = 0.5
    self._num_book = 0

  @property
  def num_book(self):
    return self._num_book

  def __iter__(self):
    return self

  def __next__(self):
    self._current_t += self._delta_t
    v = self._current_price + self.random_move()
    asks = []
    bids = []
    for n in range(1, 10, 2):
      nn = n / 2.
      asks.append([v + nn * self._tick, self._qty])
      bids.append([v - nn * self._tick, self._qty])
    self._book.snapshot(bids, asks, self._current_t)
    self._num_book += 1
    return self._book

  def random_move(self):
    return random.choice([-2, -1, -1, 0, 0, 0, 0, 1, 1, 2]) * self._tick


def run_book_generator():
  calculator = VolatilityCalculator(window_sec=30)
  book_cal = None
  for n, book in enumerate(GbmBookGenerator()):
    if book_cal is None:
      book_cal = SingleProductPriceCalculator(book.product)
    book_cal.update(book)
    calculator.update(book.timestamp, book_cal.midp)
    if calculator.is_ready():
      print(calculator.get_volatility())
    book_json = SingleProductPriceCalculator.convert_to_json(book)
    print(book_json)
    if n > 1000000:
      break


def run_from_generated_feedx(products, strategy):
  start_time = get_timestamp()
  delta_time = int(5 * 1e8)
  qty = 1e10

  t = start_time
  huobi_bids = [[10000., qty], [10000 - 0.5, qty]]
  huobi_asks = [[10001. - 0.5, qty], [10001., qty]]
  for n in range(160000):
    t += delta_time
    for prod in products:
      book = BookBuilder(prod.symbol)
      book.product = prod
      if prod.exchange == 'Huobi':
        bp = huobi_asks[0][0] / 10000
        r = random.randint(-1, 1) * bp
        if (huobi_bids[1][0] + r) < 10:
          r = bp

        huobi_bids[0][0] += r
        huobi_bids[1][0] += r
        huobi_asks[0][0] += r
        huobi_asks[1][0] += r
        book.snapshot(huobi_bids, huobi_asks, t)
      elif prod.exchange == 'Okex':
        r = random.randint(-7, 7) * 1e-4 + 1
        bids = [[huobi_bids[0][0] * r, qty], [huobi_bids[1][0] * r, qty]]
        asks = [[huobi_asks[0][0] * r, qty], [huobi_asks[1][0] * r, qty]]
        book.snapshot(bids, asks, t)
      else:
        raise ValueError()
      strategy.common_on_book(prod, book)


def run_from_generated_feed(products, strategy):
  assert len(book) == 1, products
  product = products[0]
  bookgen = RandomWalkBookGenerator()
  for book in bookgen:
    book.product = product
    strategy.common_on_book(product, book)
    if rb.num_book > 24 * 3600 * 2:
      break


def main():
  rb = RandomWalkBookGenerator()
  for book in rb:
    print(book.get_ask_array(5))
    if rb.num_book > 10:
      break


if __name__ == '__main__':
  main()
